# -*- coding: utf-8 -*-

# @Time : 2022/3/16 11:15

# @Author : Aweo
# @File : nano_predict.py

import torch
import cv2
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
from torch.autograd import Variable
from model_select import select


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 2)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        # print(x.shape)
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)

        return x


class Read():

    def __init__(self, label, model_path, model='CNN5'):
        self.label = label  #
        self.num_label = list(range(len(label)))
        self.model_path = model_path  #
        self.model = select(model)
        self.model.load_state_dict(torch.load(model_path))

    # image transforms
    def img_read(self, img):
        # img = cv2.imread(img, cv2.IMREAD_COLOR)
        img = cv2.resize(img, (32, 32))
        transform = transforms.Compose(
            [transforms.ToTensor(),
             transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
        img = Variable(torch.unsqueeze(transform(img), dim=0).float(), requires_grad=False)
        return img

    @torch.no_grad()
    def Precdict(self, img):
        img = self.img_read(img)
        pred = self.model(img)
        _, pred = torch.max(pred, 1)
        return self.label[pred]


if __name__ == '__main__':
    num_label = [0, 1]
    model_path = 'jetson_model/CNN5v2.pth'
    read = Read(label=['感染', '未感染'],
                model_path='jetson_model/CNN5v2.pth',
                model='CNN5')
    model = select('CNN5')

    # model = SEnet_train.CNN(n_class=2)
    model.load_state_dict(torch.load(model_path))
    # print(model)
    img = cv2.imread('img_data/nag_img/2000.png')
    res = read.Precdict(img)


    print(res)
    # print((pred.argmax(1) == 1).type(torch.float).sum().item())
